package com.byron.media.server.utils;

import java.util.ArrayList;
import java.util.List;

public class RtpHelper {

    public static final int MTU = 1500;		// 包的最大大小

    public static final int RTP_HEADER_LENGTH = 12;	// rtp头大小

    public static final int PAYLOAD_LENGTH = MTU - RTP_HEADER_LENGTH;	// rtp实际内容大小

    public static final int PAYLOAD_TYPE_H264 = 0x60;
//    private static int seqNum = 0;      // 包名

    public static final int PAYLOAD_TYPE_PCM = 0x28;
//    private static int seqNum = 0;      // 包名


    /**
     * 判断收到的rtp是否是单一分包
     * @param rtp
     * @return
     */
    public static boolean isSingleRtp(byte[] rtp){
        byte payloadHeader = rtp[RTP_HEADER_LENGTH];
        byte type = (byte) (payloadHeader & 0x1F);
        if(type == 28){
            return false;
        }
        return true;
    }


    public static int getSeq(byte[] rtp){
        return getShort(rtp, 2);
    }



    /**
     * 将单个rtp解包成h264
     * 忽略同步时间和同步源等RTP头包含的信息
     */
    public static byte[] unpackSingleToH264(byte[] rtp, int length){
        byte[] h264 = new byte[length - RTP_HEADER_LENGTH + 4];
        System.arraycopy(rtp, RTP_HEADER_LENGTH, h264, 4, length - RTP_HEADER_LENGTH);
        h264[3] = 0x01;
        return h264;
    }


    /**
     * 将单个rtp解包成h264
     * 忽略同步时间和同步源等RTP头包含的信息
     */
    public static byte[] unpackSingleToPCM(byte[] rtp, int length){
        byte[] pcm = new byte[length - RTP_HEADER_LENGTH];
        System.arraycopy(rtp, RTP_HEADER_LENGTH, pcm, 0, length - RTP_HEADER_LENGTH);
        return pcm;
    }


    /**
     * 将多个分包解包成h264
     */
    public static byte[] unpackMultiToH264(List<byte[]> rtps, List<Integer> lengths){
        // 统计长度
        int length = 0;
        for(Integer rtpLength : lengths){
            length += (rtpLength - RTP_HEADER_LENGTH - 2);
        }
        length += 1 + 4;
        byte[] h264 = new byte[length];

        // 还原原来的payload header
        byte[] firstRtp = rtps.get(0);
        byte payloadHeader = (byte) ((firstRtp[RTP_HEADER_LENGTH] & 0xE0) | (firstRtp[RTP_HEADER_LENGTH + 1] & 0x1F));
        h264[3] = 0x01;
        h264[4] = payloadHeader;

        int read = 5;
        for(int i = 0 ; i < rtps.size(); i++){
            byte[] rtp = rtps.get(i);
            int rtpLength = lengths.get(i);
            System.arraycopy(rtp, RTP_HEADER_LENGTH + 2, h264, read, rtpLength - RTP_HEADER_LENGTH - 2);
            read += (rtpLength - RTP_HEADER_LENGTH - 2);
        }

        assert read == length;
        return h264;
    }

    /**
     * 将多个分包解包成h264
     * @param rtps
     * @return
     */
    public static byte[] unpackMultiToPCM(List<byte[]> rtps, List<Integer> lengths){
        // 统计长度
        int length = 0;
        for(Integer rtpLength : lengths){
            length += (rtpLength - RTP_HEADER_LENGTH - 2);
        }
//        length += 1 + 4;
        byte[] pcm = new byte[length];

        // 还原原来的payload header
//        byte[] firstRtp = rtps.get(0);
//        byte payloadHeader = (byte) ((firstRtp[RTP_HEADER_LENGTH] & 0xE0) | (firstRtp[RTP_HEADER_LENGTH + 1] & 0x1F));
//        h264[3] = 0x01;
//        h264[4] = payloadHeader;

        int read = 0;
        for(int i = 0 ; i < rtps.size(); i++){
            byte[] rtp = rtps.get(i);
            int rtpLength = lengths.get(i);
            System.arraycopy(rtp, RTP_HEADER_LENGTH + 2, pcm, read, rtpLength - RTP_HEADER_LENGTH - 2);
            read += (rtpLength - RTP_HEADER_LENGTH - 2);
        }

        assert read == length;
        return pcm;
    }



    /**
     * 是否需要多个分包
     * @return
     */
    public static boolean needToMultiPackage(int length){
        if(length > PAYLOAD_LENGTH){
            return true;
        }
        return false;
    }


    /**
     * 寻找开始码
     * @param data
     * @return
     */
    public static int findStartCodeIndex(byte[] data){
        for(int i = 0; i < 10; i++){
            if(data[i] == 1){
                return i;
            }
        }
        return -1;
    }


    /**
     * 是否是分包里面最后一个分包
     * @param rtp
     * @return
     */
    public static boolean isLastMultiRtp(byte[] rtp){
        byte fuHeader = rtp[RTP_HEADER_LENGTH + 1];
        byte end = (byte) ((fuHeader & 0x40) >> 6);
        return end == 1;
    }



    /**
     * 组成多个分包
     * @return
     */
    public static List<byte[]> packageToMultiRtp(
            byte[] nalu, int offset, int length, int frameIndex, int seqNum){
        List<byte[]> rtps = new ArrayList<>();
//        int read = 0;
        //读取位置为头后
        int read = offset;
        //砍去头的nalu长度
//        length = length - (offset + 1);

        byte[] rtpHeader = new byte[RTP_HEADER_LENGTH];

        rtpHeader[0] = (byte)0x80;
        rtpHeader[1] = (byte)0x60;
        putShort(rtpHeader, (short)seqNum, 2);		// 序列号 16bit 每发一个包加1
//        putInt(rtpHeader, 3600 * frameIndex, 4);			// 时间戳  32bit 单一封包 +采样率，h264为3600; 分片封包第一个加采样率，后续不变
        putInt(rtpHeader, 0, 4);			// 时间戳  32bit 单一封包 +采样率，h264为3600; 分片封包第一个加采样率，后续不变
        putInt(rtpHeader, 0, 8);			// SSRC 同步源标识	 任意指定，标准是一个MD5算法值，未明

        // 拆分nalu header
//        byte naluHeader = nalu[0];
        byte naluHeader = nalu[read];
        byte fuIndicator = (byte) (naluHeader & 0xE0 | 0x1C);
        byte startFuHeader = (byte) (naluHeader & 0x1F | 0x80);
        byte middleFuHeader = (byte)(naluHeader & 0x1F);
        byte endFuHeader = (byte)(naluHeader & 0x1F | 0x40);
        read += 1;

        while(read < length){

            int maxPayloadLength = PAYLOAD_LENGTH - 2;
            int rtpPayloadLength = maxPayloadLength > (length - read) ? (length - read) : maxPayloadLength;
            int rtpLength = rtpPayloadLength + RTP_HEADER_LENGTH + 2;

            byte[] rtp = new byte[rtpLength];

            System.arraycopy(nalu, read, rtp, RTP_HEADER_LENGTH + 2, rtpPayloadLength);
            read += rtpPayloadLength;

            // 判断是不是最后一个包
            if(read == length){

                // 包名加一
                seqNum ++;
                putShort(rtpHeader, (short)seqNum, 2);		// 序列号 16bit 每发一个包加1
//                // 修改M值
//                byte mPT = rtpHeader[1];
//                mPT = (byte) (mPT | 0x80);
//                rtpHeader [1] = mPT;
//                System.arraycopy(rtpHeader, 0, rtp, 0, RTP_HEADER_LENGTH);
                rtpHeader[1] = (byte)0xE0;
                System.arraycopy(rtpHeader, 0, rtp, 0, RTP_HEADER_LENGTH);

                // 修改FuHeader
                rtp[0 + RTP_HEADER_LENGTH] = fuIndicator;
                rtp[1 + RTP_HEADER_LENGTH] = endFuHeader;

                rtps.add(rtp);

                // 判断开头
            } else if(read - 4 == rtpPayloadLength + 1){

                seqNum ++;
                putShort(rtpHeader, (short)seqNum, 2);		// 序列号 16bit 每发一个包加1
                System.arraycopy(rtpHeader, 0, rtp, 0, RTP_HEADER_LENGTH);

                rtp[0 + RTP_HEADER_LENGTH] = fuIndicator;
                rtp[1 + RTP_HEADER_LENGTH] = startFuHeader;

                rtps.add(rtp);

            } else {
                seqNum ++;
                putShort(rtpHeader, (short)seqNum, 2);		// 序列号 16bit 每发一个包加1
                System.arraycopy(rtpHeader, 0, rtp, 0, RTP_HEADER_LENGTH);

                rtp[0 + RTP_HEADER_LENGTH] = fuIndicator;
                rtp[1 + RTP_HEADER_LENGTH] = middleFuHeader;

                rtps.add(rtp);
            }
        }


        return rtps;
    }


    /**
     * 获取PayloadType
     * @param rtp
     * @return
     */
    public static int getPayloadType(byte[] rtp){
        byte type = rtp[1];
        return type & 0x7F;
    }


    /**
     * 组成多个分包
     * @return
     */
    public static List<byte[]> packagePcmToMultiRtp(
            byte[] nalu, int offset, int length, int frameIndex, int seqNum){
        List<byte[]> rtps = new ArrayList<>();
//        int read = 0;
        //读取位置为头后
        int read = 0;
        //砍去头的nalu长度
//        length = length - (offset + 1);

        byte[] rtpHeader = new byte[RTP_HEADER_LENGTH];

        rtpHeader[0] = (byte)0x80;
        rtpHeader[1] = (byte)0x28;              // 自定义PCM 的PT 为40
        putShort(rtpHeader, (short)seqNum, 2);		// 序列号 16bit 每发一个包加1
        putInt(rtpHeader, 0, 4);			// 时间戳  32bit 单一封包 +采样率，h264为3600; 分片封包第一个加采样率，后续不变
        putInt(rtpHeader, 0, 8);			// SSRC 同步源标识	 任意指定，标准是一个MD5算法值，未明

        // 拆分nalu header
//        byte naluHeader = nalu[0];
        byte naluHeader = nalu[read];
        byte fuIndicator = (byte) (naluHeader & 0xE0 | 0x1C);
        byte startFuHeader = (byte) (naluHeader & 0x1F | 0x80);
        byte middleFuHeader = (byte)(naluHeader & 0x1F);
        byte endFuHeader = (byte)(naluHeader & 0x1F | 0x40);
        read += 1;

        while(read < length){

            int maxPayloadLength = PAYLOAD_LENGTH - 2;
            int rtpPayloadLength = maxPayloadLength > (length - read) ? (length - read) : maxPayloadLength;
            int rtpLength = rtpPayloadLength + RTP_HEADER_LENGTH + 2;

            byte[] rtp = new byte[rtpLength];

            System.arraycopy(nalu, read, rtp, RTP_HEADER_LENGTH + 2, rtpPayloadLength);
            read += rtpPayloadLength;

            // 判断是不是最后一个包
            if(read == length){

                // 包名加一
                seqNum ++;
                putShort(rtpHeader, (short)seqNum, 2);		// 序列号 16bit 每发一个包加1
//                // 修改M值
//                byte mPT = rtpHeader[1];
//                mPT = (byte) (mPT | 0x80);
//                rtpHeader [1] = mPT;
//                System.arraycopy(rtpHeader, 0, rtp, 0, RTP_HEADER_LENGTH);
                rtpHeader[1] = (byte)0xA8;
                System.arraycopy(rtpHeader, 0, rtp, 0, RTP_HEADER_LENGTH);

                // 修改FuHeader
                rtp[0 + RTP_HEADER_LENGTH] = fuIndicator;
                rtp[1 + RTP_HEADER_LENGTH] = endFuHeader;

                rtps.add(rtp);

                // 判断开头
            } else if(read - 4 == rtpPayloadLength + 1){

                seqNum ++;
                putShort(rtpHeader, (short)seqNum, 2);		// 序列号 16bit 每发一个包加1
                System.arraycopy(rtpHeader, 0, rtp, 0, RTP_HEADER_LENGTH);

                rtp[0 + RTP_HEADER_LENGTH] = fuIndicator;
                rtp[1 + RTP_HEADER_LENGTH] = startFuHeader;

                rtps.add(rtp);

            } else {
                seqNum ++;
                putShort(rtpHeader, (short)seqNum, 2);		// 序列号 16bit 每发一个包加1
                System.arraycopy(rtpHeader, 0, rtp, 0, RTP_HEADER_LENGTH);

                rtp[0 + RTP_HEADER_LENGTH] = fuIndicator;
                rtp[1 + RTP_HEADER_LENGTH] = middleFuHeader;

                rtps.add(rtp);
            }
        }


        return rtps;
    }

    /**
     * 组成单个rtp包
     * @param nalu
     * @param length
     * @return
     */
    public static byte[] packageToRtp(byte[] nalu, int offset, int length, int frameIndex, int seqNum){
        //砍掉头后的length
        byte[] rtp = new byte[length + RTP_HEADER_LENGTH - offset];
        rtp[0] = (byte)0x80;
        rtp[1] = (byte)0xE0;
        seqNum++;
        putShort(rtp, (short)seqNum, 2);		// 序列号 16bit 每发一个包加1
        putInt(rtp, 0, 4);
//        putInt(rtp, 3600 * frameIndex, 4);			// 时间戳  32bit 单一封包 +采样率，h264为3600; 分片封包第一个加采样率，后续不变
        putInt(rtp, 0, 8);			// SSRC 同步源标识	 任意指定，标准是一个MD5算法值，未明
        System.arraycopy(nalu, offset, rtp, RTP_HEADER_LENGTH, length - offset);
        return rtp;
    }

    /**
     * 组成单个rtp包
     * @param nalu
     * @param length
     * @return
     */
    public static byte[] packagePcmToRtp(byte[] nalu, int offset, int length, int frameIndex, int seqNum){
        //砍掉头后的length
        byte[] rtp = new byte[length + RTP_HEADER_LENGTH - offset];
        rtp[0] = (byte)0x80;
        rtp[1] = (byte)0x28;
        seqNum++;
        putShort(rtp, (short)seqNum, 2);		// 序列号 16bit 每发一个包加1
        putInt(rtp, 0, 4);
//        putInt(rtp, 3600 * frameIndex, 4);			// 时间戳  32bit 单一封包 +采样率，h264为3600; 分片封包第一个加采样率，后续不变
        putInt(rtp, 0, 8);			// SSRC 同步源标识	 任意指定，标准是一个MD5算法值，未明
        System.arraycopy(nalu, offset, rtp, RTP_HEADER_LENGTH, length - offset);
        return rtp;
    }

    /**
     * 转换short为byte
     *
     * @param b
     * @param s 需要转换的short
     * @param index
     */
    private static void putShort(byte b[], short s, int index) {
        b[index + 1] = (byte) (s >> 0);
        b[index + 0] = (byte) (s >> 8);
    }


    /**
     * 通过byte数组取到short
     * 由于平台兼容性,高位在后4位,低位在前4位
     * @param b
     * @param index 第几位开始取
     * @return
     */
    public static short getShort(byte[] b, int index) {
        return (short) (((b[index + 1] & 0xff ) | b[index + 0]<<8 & 0xff00));
    }


    /**
     * 转换Int为byte
     *
     * @param b
     * @param s 需要转换的Int
     * @param index
     */
    private static void putInt(byte b[] , int s ,int index){
        b[index + 3] = (byte) (s >> 0);
        b[index + 2] = (byte) (s >> 8);
        b[index + 1] = (byte) (s >> 16);
        b[index + 0] = (byte) (s >> 24);
    }

    /**
     * 由于I帧是和SPS PPS 一起发送, 所以只需要判断SPS 的类型就行
     * @param frameData
     * @return
     */
    public static boolean checkH264IsIFrame(byte[] frameData) {

        if(frameData.length > 4){
            if((frameData[4] & 0x1f) == 7){
                return true;
            } else if((frameData[4] & 0x1f) == 9){
                return true;
            }
        }
        return false;
    }

}
